import numpy as np
from pylab import *
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.image as mpimg
from Interpolation import *
from Sector_Parameters import *

font = {'family' : 'monospace', 'weight' : 'bold', 'size'   : 40}
mpl.rc('font', **font)  # pass in the font dict as kwargs
mpl.rcParams['axes.linewidth'] = 10

Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,N_p, Nx_aug,sigma_pnts=circular_parameters()
DeltaT, Tsim, Nsim=time_parameters()
Nt_new = 85
Nr_new = 146
def surface_parameters(d,par,str1, filename):
    P=[]
    for j in range(Nt):
        c=d[int(((Nt+j)-Nt)*Nr):int(((Nt+j+1)-Nt)*Nr)] 
        P.append(c)
    PP=np.array(P)
    z=PP
    z_new=new_values(z)
    z_new_ = new_values_radial(z_new)
    Z_new = z_new_.squeeze()
    Z_new = np.transpose(Z_new)
    fig = plt.figure()
    r_new=np.linspace(0,290,Nr_new)
    th_new = np.linspace(0, np.pi/2, Nt_new)
    R, TH = meshgrid(r_new, th_new)
    
    fig=plt.figure(figsize=(15,15))
    ax=Axes3D(fig)
    m=plt.subplot(projection="polar")
    m.set_xticklabels([])
    xlim([0,pi/2])# Creating a sector from the  
    m.set_yticklabels([])
    plt.pcolormesh(TH,R,Z_new,cmap='turbo',shading='flat')
    plt.colorbar(label=str1, shrink=0.8)
    plt.clim(Z_new.min(),Z_new.max())
    fig.tight_layout()
    return Z_new

def surface_parameters_orig(d,par,str1, filename):
    Z_new = d
    fig=plt.figure(figsize=(25,25))
    ax=Axes3D(fig)
    r_new = np.linspace(0, 290, Nr_new)
    th_new = np.linspace(0, np.pi/2, Nt_new)
    R, TH = meshgrid(r_new, th_new)
    m=plt.subplot(projection="polar")
    m.set_xticklabels([])
    m.grid(linewidth=0, color='k')
    xlim([0,np.pi/2])
    m.set_yticklabels([])
    colorList = ['#641E16','#512E5F', '#154360','#0E6251','#145A32']
    cmap = matplotlib.colors.ListedColormap(colorList)
    bounds = [0.,1.,2.,3.,4.]
    ticks = [0.5,1.5,2.5,3.5, 4.5]
    p = plt.pcolormesh(TH,R,Z_new,cmap=cmap, vmin = 0, vmax=5, shading='flat', linewidth=0, linestyle='None')
    cbar = fig.colorbar(p, aspect = 20, ticks=ticks, shrink=0.75)
    cbar.set_ticklabels(['MZ$_1$','MZ$_2$','MZ$_3$', 'MZ$_4$','MZ$_5$'])
    cbar.set_label(label='Legend', weight='bold', family='monospace', size=40, labelpad= 14)
    fig.tight_layout()
    plt.savefig('./results/MZs_Q1.jpg')
    plt.show()